import os
import tensorflow as tf
import pandas as pd
from sklearn.model_selection import train_test_split


# 获取切分好的数据集
def get_datasets():
    data = pd.read_csv('../archive/data.csv')
    # print(data.head())
    print(data.shape)

    # 设置target和input_data
    X, y = data.drop('Bankrupt?', axis=1), data['Bankrupt?']
    print(X.shape, y.shape)

    # 分割训练集和测试集(设置训练集为全部数据的30%)
    X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.3)
    return X_train, y_train, X_test, y_test
